package com.lw.leetcode.tree.c;

/**
 * Created with IntelliJ IDEA.
 * 710. 黑名单中的随机数
 * @author liw
 * @version 1.0
 * @date 2022/6/23 17:14
 */
public class Solution {

    private int all;
    private Node root;

    public Solution(int n, int[] blacklist) {
        root = new Node(0, n - 1);
        for (int v : blacklist) {
            add(root, v);
        }
        all = n - blacklist.length;
    }

    public int pick() {
        return find(root, (int) (Math.random() * all));
    }

    private int find(Node node, int val) {
        Node left = node.left;
        Node right = node.right;
        int st = node.st;
        int end = node.end;
        int m = ((end - st) >> 1) + 1;
        int l = left == null ? m : left.end - left.st + 1 - left.count;
        if (val < l) {
            if (left == null) {
                return st + val;
            }
            return find(node.left, val);
        }
        if (right == null) {
            return st + m + val - l;
        }
        return find(node.right, val - l);
    }

    private void add(Node node, int val) {
        int st = node.st;
        int end = node.end;
        node.count++;
        if (st == end) {
            return;
        }
        int m = st + ((end - st) >> 1);
        if (val <= m) {
            if (node.left == null) {
                node.left = new Node(st, m);
            }
            add(node.left, val);
        } else {
            if (node.right == null) {
                node.right = new Node(m + 1, end);
            }
            add(node.right, val);
        }
    }

    private static class Node {
        private int st;
        private int end;
        private int count;
        private Node left;
        private Node right;

        private Node(int st, int end) {
            this.st = st;
            this.end = end;
        }
    }

}
